from functools import wraps
import asyncio
from redis import asyncio as aredis
from .cache import get_redis, RedisLock



# 解析异步参数
async def _resolve(arg):
    return await arg if asyncio.iscoroutine(arg) else arg

def lock_fn(
    redis_client: aredis.Redis,
    get_lock_key,
    lock_timeout: int = 10,
    acquire_timeout: int = 3,
    nowait: bool = False,
):
    def decorator(fn):
        @wraps(fn)
        async def wrapper(*args, **kwargs):
            args_values = await asyncio.gather(*[_resolve(arg) for arg in args])
            kwargs_values = {k: await _resolve(v) for k, v in kwargs.items()}
            key = get_lock_key(*args_values, **kwargs_values)
            lock_key = f"lock:{key}"

            async with RedisLock(
                redis_client, lock_key, lock_timeout, acquire_timeout, nowait=nowait
            ):
                return await fn(*args, **kwargs)

        return wrapper

    return decorator
